import asyncio
import json
import time

import httpx
from fastapi import FastAPI, UploadFile, Form, HTTPException, File
from fastapi.responses import JSONResponse
import uvicorn

# 配置
SERVER_HTTP_URL = "http://192.168.31.207:5000"  # 目标FastAPI服务器地址
RELAY_PORT = 5000  # 中继服务器端口
MAX_UPLOAD_SPEED = 8 * 1024 * 1024  # 最大上传速度 8Mb/s

relay = FastAPI()


@relay.post("/upload/{link_id}")
async def relay_upload(link_id: str, file: UploadFile = File(...), metadata: str = Form(...)):
    async def read_file_with_limit(file: UploadFile, max_speed: int) -> bytes:
        file_content = b''
        start_time = time.time()  # 记录开始时间
        total_read = 0  # 记录总共读取的字节数

        while True:
            # 计算本次读取的最大字节数
            chunk_size = min(max_speed, max_speed - total_read)
            if chunk_size <= 0:
                break

            # 读取数据块
            chunk = await file.read(chunk_size)
            if not chunk:
                break

            file_content += chunk
            total_read += len(chunk)

            # 计算已经过去的时间
            elapsed_time = time.time() - start_time
            expected_time = total_read / max_speed  # 预期的时间

            # 如果实际时间小于预期时间，则需要等待
            if elapsed_time < expected_time:
                await asyncio.sleep(expected_time - elapsed_time)

        return file_content

    try:
        file_metadata = json.loads(metadata)
        is_accept = file_metadata['chunk_index'] % 1 == 0
        if is_accept:
            file_content = await read_file_with_limit(file, MAX_UPLOAD_SPEED)
            files = {'file': (file.filename, file_content)}
            data = {'metadata': metadata}

            async with httpx.AsyncClient() as client:
                response = await client.post(
                    f"{SERVER_HTTP_URL}/upload/{link_id}",
                    files=files,
                    data=data
                )

            if response.headers.get("content-type") == "application/json":
                return JSONResponse(
                    content=response.json(),
                    status_code=response.status_code
                )
            else:
                return JSONResponse(
                    content={"error": "Unexpected response from server"},
                    status_code=500
                )
        else:
            return JSONResponse(content={"message": "Chunk uploaded unsuccessfully", "chunk_index": file_metadata['chunk_index']},
                                status_code=400)

    except httpx.HTTPError as e:
        raise HTTPException(status_code=500, detail=f"Failed to forward request: {str(e)}")
    except json.JSONDecodeError as e:
        raise HTTPException(status_code=400, detail=f"Invalid metadata: {str(e)}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Unexpected error: {str(e)}")


async def start_server():
    """启动中继服务器"""
    config = uvicorn.Config(
        relay,
        host="0.0.0.0",
        port=RELAY_PORT,
        log_level="warning",
        timeout_keep_alive=60
    )
    server = uvicorn.Server(config)
    await server.serve()


if __name__ == "__main__":
    asyncio.run(start_server())